import "dotenv/config";
import express from "express";
import { connectDB } from "./config/db.mjs";
import http from "http";
import { WebSocketServer } from "ws";
import { ZhipuSDK } from "./utils/ZhipuSDK.mjs";

// 导入路由
import schoolRoutes from "./routes/school/index.mjs";
import provinceRoutes from "./routes/province/index.mjs";
import aiRoutes from "./routes/ai/index.mjs";
import statisticsRoutes from "./routes/statistics/index.mjs";

// 导入swagger
import swaggerInstall from "./swagger/index.mjs";
import cors from "cors";

const app = express();
const server = http.createServer(app);

// 创建WebSocket服务器
const wss = new WebSocketServer({ server });

// 存储所有连接的WebSocket客户端
const clients = new Set();

// WebSocket连接处理
wss.on("connection", (ws) => {
  clients.add(ws);
  console.log("新的WebSocket连接已建立");

  ws.on("message", async (message) => {
    try {
      const data = JSON.parse(message);
      // console.log("收到消息:", data);

      if (data.type === "chat") {
        // 处理聊天请求
        handleChatRequest(ws, data.messages);
      }
    } catch (error) {
      console.error("WebSocket message error:", error);
      ws.send(
        JSON.stringify({
          type: "error",
          content: "处理请求时发生错误",
        })
      );
    }
  });

  ws.on("close", () => {
    clients.delete(ws);
    console.log("WebSocket连接已关闭");
  });
});

// 处理聊天请求
async function handleChatRequest(ws, messages) {
  try {
    // 发送开始处理的消息
    ws.send(
      JSON.stringify({
        type: "status",
        content: "正在处理您的请求...",
      })
    );

    // 创建智谱AI实例
    const ai = new ZhipuSDK({
      apiKey: process.env.ZHIPU_API_KEY,
    });

    // 一次性获取完整回复
    const response = await ai.createCompletions({
      model: "glm-4-flash-250414",
      tools: [
        {
          type: "retrieval",
          retrieval: {
            knowledge_id: "1907314965451272192",
            prompt_template: `
              请基于知识库中的院校信息进行分析和推荐。
              如果用户询问具体院校信息，请仅推荐知识库中实际存在的院校，并根据知识库中院校的介绍，给出推荐理由。
              如果用户询问到录取分数线、GDP、就业率、收入增长、生活成本、教育资源等内容，请从知识库中存在的数据中获取数据

              如果无法在知识库中找到相关院校，请明确告知用户,。
            `,
          },
        },
      ],
      messages: [
        {
          role: "system",
          content: `
          你是一个高校数据智能分析师，名为"AI小助手"。你的核心能力包括：
          
          【身份定义】
            1. 专业定位：高校数据智能分析师，熟悉2018-2024年全国高校招生政策
            2. 服务范围：提供院校对比、专业解读、往年录取分数查询、分数线趋势分析服务、志愿推荐功能

          【重要规则】
            1. 院校推荐规则：
            - 只能推荐知识库中实际存在的院校
            - 每次推荐院校时，必须先查询知识库确认院校信息
            - 如果要推荐的院校不在知识库中，应该寻找知识库中类似的替代院校
    
            2. 回答格式要求：
            - 推荐院校时，使用标准的学校全称
            - 确保推荐的院校数量适中
            - 每所推荐院校都应该附带推荐理由
            - 介绍学校时根据知识库中信息介绍
            - 如果用户询问了数据请根据知识库中数据进行回答，若是没有询问数据，只需根据数据进行分析，不需要向用户展示数据

            3. 其他相关数据： 
            - 录取分数线、GDP、就业率、收入增长、生活成本、教育资源等内容，必须从知识库中获取数据。
            如果知识库中不存在相关数据，应明确告知用户，并尝试提供其他相关知识库中的数据或信息作为替代。  
            
            【智能问答流程】
            1. 解析用户输入的关键维度：
            - 分数段（如550-600分）
            - 地域偏好（如江浙沪地区）
            - 专业兴趣（如工科类）
            - 相关数据：录取分数线、GDP、就业率、收入增长、生活成本、教育资源
    
            2. 分析步骤：
            - 先在知识库中筛选符合条件的院校
            - 结合用户需求进行匹配分析
            - 确保推荐的院校都在知识库范围内
            - 确保数据支持的来源是知识库
            - 提供详细的推荐理由和数据支持
    
            3. 输出规范：
            - 确保学校名称准确无误
            - 包含推荐理由和相关数据支持
            - 如果知识库中缺乏特定数据，提供替代信息并说明原因

            ## Knowledge
            请根据以下知识库中的数据，给出推荐理由和相关数据支持，
            数据包含全国高校排行、分类、类别、2018-2024年录取分数线、GDP、就业率、收入增长、生活成本、教育资源
            """{{knowledge}}"""

          `,
        },
        ...messages,
      ],
      stream: false,
    });

    // 获取完整回复内容
    const fullResponse = response.choices[0].message.content;

    // 发送完整回复
    ws.send(
      JSON.stringify({
        type: "message",
        content: fullResponse,
      })
    );

    // 提取学校、省份实体信息
    const summaryResponse = await ai.createCompletions({
      model: "glm-4-flash-250414",
      messages: [
        {
          role: "system",
          content:
            '你是一个高考志愿填报助手，需要从对话中精确提取推荐的大学院校和省份或直辖市，并以严格JSON格式返回。院校必须来自知识库，格式示例：{"recommended_schools": ["清华大学","北京大学"], "recommended_provinces": ["省份1","省份2"]}',
        },
        {
          role: "user",
          content: fullResponse,
        },
      ],
      stream: false,
      response_format: {
        type: "json_object",
      },
      tools: [
        {
          type: "retrieval",
          retrieval: {
            knowledge_id: "1907314965451272192",
            prompt_template: `
            # 任务要求
            1. 严格从知识库colleges-geo-info.csv包含的院校中提取，同时返回该院校所在省份或直辖市
            2. 只返回规范的JSON格式：{"recommended_schools": ["院校1","院校2"], "recommended_provinces": ["省份1(如：北京市)","省份2(如：广东省)"]}
            3. 若对话中没有有效院校或院校不在知识库中，返回：{"recommended_schools": [], "recommended_provinces": []}
    
            ## Knowledge
            """{{knowledge}}"""
    
            ## Workflow
            1、读取用户问题："{{question}}"
            2、请理解并记住以下内容：
            """{{knowledge}}"""
            `,
          },
        },
      ],
    });

    // 处理推荐结果
    try {
      const recommendationsContent = summaryResponse.choices[0].message.content;
      console.log("推荐结果：", recommendationsContent);
      let recommended_schools;
      let recommended_provinces;

      try {
        recommended_schools = JSON.parse(
          recommendationsContent
        ).recommended_schools;
        recommended_provinces = JSON.parse(
          recommendationsContent
        ).recommended_provinces;
      } catch {
        // 如果解析失败，尝试从文本中提取学校名称和省份
        // 修改正则表达式，匹配学校名称
        recommended_schools =
          recommendationsContent
            .match(/["']([^"']+(?:大学|学院|师范|科技)[^"']*)["']/g)
            ?.map((s) => s.replace(/["']/g, ""))
            ?.filter((name) => name !== "recommended_schools") || [];

        // 优化省份提取逻辑
        const provincePattern =
          /["']([^"']+(?:省|市|自治区|特别行政区|北京|上海|天津|重庆)[^"']*)["']/g;
        recommended_provinces = [];
        let match;
        while (
          (match = provincePattern.exec(recommendationsContent)) !== null
        ) {
          const province = match[1].replace(/["']/g, "");
          // 标准化省份名称
          const normalizedProvince = province
            .replace(/(?:省|市|自治区|特别行政区)$/, "")
            .replace(/^北京市$/, "北京")
            .replace(/^上海市$/, "上海")
            .replace(/^天津市$/, "天津")
            .replace(/^重庆市$/, "重庆");
          if (!recommended_provinces.includes(normalizedProvince)) {
            recommended_provinces.push(normalizedProvince);
          }
        }
      }

      // 确保推荐学校是一个数组
      if (!Array.isArray(recommended_schools)) {
        recommended_schools = [];
      }

      if (!Array.isArray(recommended_provinces)) {
        recommended_provinces = [];
      }

      // 只有当有推荐学校时才发送推荐消息
      if (recommended_schools.length > 0) {
        ws.send(
          JSON.stringify({
            type: "recommendations-schools",
            content: JSON.stringify(recommended_schools),
          })
        );
      }
      if (recommended_provinces.length > 0) {
        ws.send(
          JSON.stringify({
            type: "recommendations-provinces",
            content: JSON.stringify(recommended_provinces),
          })
        );
      }
    } catch (error) {
      console.error("处理推荐结果时出错:", error);
      // 推荐处理出错时不发送推荐消息
    }

    // 发送完成消息
    ws.send(
      JSON.stringify({
        type: "done",
        content: "处理完成",
      })
    );
  } catch (error) {
    console.error("Chat error:", error);
    ws.send(
      JSON.stringify({
        type: "error",
        content: "处理请求时发生错误",
      })
    );
  }
}

swaggerInstall(app);

// 连接MongoDB
connectDB();

// 中间件
app.use(express.json());
app.use(cors());

// 路由
app.use("/api/schools", schoolRoutes);
app.use("/api/provinces", provinceRoutes);
app.use("/api/ai", aiRoutes);
app.use("/api/statistics", statisticsRoutes);

// 错误处理中间件
app.use((err, req, res, next) => {
  console.error(err.stack);
  res.status(500).json({
    message: err.message,
    stack: process.env.NODE_ENV === "development" ? err.stack : undefined,
  });
});

// 启动服务器
const PORT = process.env.PORT || 3000;
server.listen(PORT, () => {
  console.log(`服务器运行在端口${PORT}`);
});
